#!/usr/bin/env python3
#coding: utf-8

import time
import re
import sys
import stat
import paramiko
from execute_command import ExecuteCommand
from execute_command import decode_command_result
from execute_log import server_logger

def auto_reconnect(func):
    def _wapper(self, *args, **kwargs):
        retries = 0
        while retries < self.max_retries:
            try:
                if not self.is_connected() or self.execute_cnt > 50:
                    self.execute_cnt = 0
                    self.connect()
                return func(self, *args, **kwargs)
            except paramiko.SSHException as e:
                server_logger.error(f'SSH connection exception: {e}')
                self.disconnected()
                server_logger.error(f'Build a new SSH connection and attempt to reconnect')
                retries += 1
                time.sleep(self.retry_interval)
            except Exception as e:
                server_logger.error(f'An unknown error was encountered while auto_reconnect: {e}')
    return _wapper

def sftp_try(func):
    def _wapper(self, *args, **kwargs):
        try:
            self.sftp_start_time = time.time()
            self.progress_cnt = 0
            return func(self, *args, **kwargs)
        except FileNotFoundError as e:
            server_logger.error(f'file does not exist: {e.filename}')
        except Exception as e:
            server_logger.error(f'An unknown error was encountered while _wapper')
    return _wapper

def decode_channel(channel: paramiko.ChannelFile) -> str:
    return decode_command_result(channel.read())

class SSHClient():
    ssh: paramiko.SSHClient
    sftp: paramiko.SFTPClient
    username: str
    password: str
    host_ip: str
    host_port: int
    timeout: int
    max_retries: int
    retry_interval: int
    get_pty: bool
    sftp_start_time: float
    progress_cnt: int
    execute_cnt: int

    def __init__(self,
                 username,
                 password,
                 host_ip,
                 host_port=22,
                 timeout=30,
                 max_retries=5,
                 retry_interval=5,
                 get_pty=False) -> None:
        if isinstance(server_logger, str):
            paramiko.util.log_to_file(server_logger)
        self.username = username
        self.password = password
        self.host_ip = host_ip
        self.host_port = host_port
        self.timeout = timeout
        self.max_retries = max_retries
        self.retry_interval = retry_interval
        self.get_pty = get_pty
        self.execute_cnt = 0

    def connect(self):
        retries = 0
        while retries < self.max_retries:
            try:
                server_logger.info(f'start SSH connection: ssh {self.username}@{self.host_ip}')
                self.ssh = paramiko.SSHClient()
                self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                self.ssh.connect(self.host_ip, self.host_port, self.username, self.password, timeout=self.timeout)

                self.sftp = self.ssh.open_sftp()
                self.sftp.sock.settimeout(self.timeout)
                server_logger.info(f'SSH connection successful')
                return
            except paramiko.AuthenticationException as e:
                server_logger.error(f'SSH authentication failed: {e}')
                raise
            except paramiko.SSHException as e:
                server_logger.error(f'SSH connection exception: {e}')
                retries += 1
                time.sleep(self.retry_interval)
            except Exception as e:
                server_logger.error(f'Unknown error occurred during SSH connection: {e}')
                raise
        raise Exception("SSH connection failed, reconnection attempts: {retries}")

    def is_connected(self):
        if (not hasattr(self, 'sftp')) or (self.sftp is None) or (self.sftp.get_channel() is None) or (self.sftp.get_channel().closed):
            return False
        if (not hasattr(self, 'ssh')) or  (self.ssh is None) or (self.ssh.get_transport() is None) or (not self.ssh.get_transport().is_active()):
            return False
        return True

    def disconnected(self):
        if hasattr(self, 'sftp') and self.sftp is not None:
            self.sftp.close()
        if hasattr(self, 'ssh') and self.ssh is not None:
            self.ssh.close()

    def print_progress(self, transferred, total):
        if self.progress_cnt < 100:
            self.progress_cnt += 1
            return
        self.progress_cnt = 0
        mb_transferred = transferred / (1024 * 1024)
        mb_total = total / (1024 * 1024)
        duration = time.time() - self.sftp_start_time
        server_logger.info(f'\rTransferred: {mb_transferred:.2f}/{mb_total:.2f} MB Speed: {mb_transferred/duration:.2f} MB/s')

    @auto_reconnect
    @sftp_try
    def send_file(self, local_file, remote_file) -> bool:
        server_logger.info(f'send file {local_file} -> {remote_file}')
        try:
            self.sftp.put(local_file, remote_file, callback=self.print_progress)
        except Exception as e:
            server_logger.error(f'send_file error: (e)')
            return False
        return True

    @auto_reconnect
    def recv_file(self, remote_file, local_file) -> None:
        server_logger.info(f'recv file {remote_file} -> {local_file}')
        self.sftp.get(remote_file, local_file, callback=self.print_progress)

    @auto_reconnect
    def mkdir_if_not_exists(self, remote_path) -> bool:
        server_logger.info(f'mkdir if not exists {remote_path}')
        path_back = self.sftp.normalize('.')
        paths = remote_path.split('\\')
        for cur_path in paths:
            try:
                self.sftp.chdir(cur_path)
            except FileNotFoundError:
                self.sftp.mkdir(cur_path)
                self.sftp.chdir(cur_path)
            except Exception as e:
                server_logger.error(f'mkdir_if_not_exists error: {e}')
                return False
        self.sftp.chdir(path_back)
        return True

    @auto_reconnect
    def clear_directory(self, remote_path) -> None:
        server_logger.info(f'clear directory {remote_path}')
        try:
            self.sftp.stat(remote_path)
        except IOError:
            return False
        for item in self.sftp.listdir_attr(remote_path):
            item_path = remote_path + '\\' + item.filename
            if stat.S_ISDIR(item.st_mode):
                clear_directory(self, item_path)
                self.sftp.rmdir(item_path)
            else:
                self.sftp.remove(item_path)
        return True

    @auto_reconnect
    def execute(self, command: list[str], timeout: int = 30, cwd: None | str = None) -> [bool, str]:
        cmd_str = ''.join(command)
        server_logger.info(f'execute command: {cmd_str}')
        # Each call to exec_command creates a new channel. 
        # When SSH connects to Win10 and frequently calls the exec_command interface, if the number of channels exceeds 120, a channel closed error will be triggered. 
        # The reason is unknown.
        self.execute_cnt = self.execute_cnt + 1
        try:
            stdin, stdout, stderr = self.ssh.exec_command(cmd_str, timeout=timeout, get_pty=self.get_pty)
            errstr = decode_channel(stderr)
            if errstr != '':
                server_logger.error(f'Remote command execution failed: {errstr}')
                return [False, errstr]
            else:
                output = decode_channel(stdout)
                return [True, output]
        except TimeoutError as e:
            err_msg = f'Remote command execution timed out: {e}'
            server_logger.info(err_msg)
            stdout.close()
            return [False, err_msg]
        except Exception as e:
            return [False, f'Remote command execution failed: {e}']
